from sentence_transformers import SentenceTransformer

"""
sentence-transformers 库专注于易用的句子嵌入表示，主要用于稠密检索。它底层基于 transformers，并提供了更简单的 API 来获取句子的稠密向量
"""
model = SentenceTransformer('D:/ideaSpace/MyPython/models/bge-m3')
sentences = ["Hello, world!", "你好，世界！"]
dense_embeddings = model.encode(sentences)
print(dense_embeddings.shape) # 输出： (2, 1024)


"""使用 Transformers 库（通用性强）
直接使用 transformers 的 AutoModel 默认只会加载模型的稠密（Dense）表示部分，可能无法直接获得稀疏（Sparse）和多向量（ColBERT）表示
"""
from transformers import AutoTokenizer, AutoModel
import torch

# 加载模型和分词器
model_name = "D:/ideaSpace/MyPython/models/bge-m3"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name, torch_dtype=torch.float16)

# 编码句子
sentences = ["Hello, world!", "你好，世界！"]
inputs = tokenizer(sentences, padding=True, truncation=True, return_tensors='pt')
with torch.no_grad():
    outputs = model(**inputs)
    # 取 [CLS] token 的输出作为句子的稠密向量
    dense_embeddings = outputs.last_hidden_state[:, 0]

print(dense_embeddings.shape) # 输出： torch.Size([2, 1024])



"""
使用官方推荐的 FlagEmbedding 库（最方便，功能最全）
这是最直接的方式，能完美支持 BGE-M3 的所有特性（稠密检索、稀疏检索、多向量检索）
"""
from FlagEmbedding import BGEM3FlagModel

model = BGEM3FlagModel("D:/ideaSpace/MyPython/models/bge-m3", use_fp16=True)
passage = ["猢狲施展烈焰拳，击退妖怪；随后开启金刚体，抵挡神兵攻击。"]

# 编码文本，获取稀疏嵌入和密集嵌入
passage_embeddings = model.encode(
    passage,
    return_sparse=True,     # 返回稀疏嵌入
    return_dense=True,      # 返回密集嵌入
    return_colbert_vecs=True  # 返回多向量嵌入
)
# 分别提取稀疏嵌入、密集嵌入和多向量嵌入
dense_vecs = passage_embeddings["dense_vecs"]
sparse_vecs = passage_embeddings["lexical_weights"]
colbert_vecs = passage_embeddings["colbert_vecs"]
# 展示稀疏嵌入和密集嵌入的示例
print("密集嵌入维度:", dense_vecs[0].shape)
print("密集嵌入前10维:", dense_vecs[0][:10])  # 仅显示前10维

print("稀疏嵌入总长度:", len(sparse_vecs[0]))
print("稀疏嵌入前10个非零值:", list(sparse_vecs[0].items())[:10])  # 仅显示前10个非零值

print("多向量嵌入维度:", colbert_vecs[0].shape)
print("多向量嵌入前2个:", colbert_vecs[0][:2])  # 仅显示前2个多向量嵌入